package org.example.myleet.unionfind;

import org.example.myleet.Utils.Assert;
import org.example.myleet.Utils.JsonUtils;
import org.example.myleet.Utils.Parser;

import java.util.*;

public class UnionFindSolution {

    static class UnionFind {
        int[] parentMap;

        public UnionFind(int n) {
            parentMap = new int[n];
            for (int i = 0; i < n; ++i) {
                parentMap[i] = i;
            }
        }

        public void union(int x, int y) {
            int px = find(x);
            int py = find(y);
            parentMap[px] = py;
        }

        public int find(int x) {
            if (parentMap[x] == x) {
                return x;
            }
            return parentMap[x] = find(parentMap[x]);
        }
    }
    /**
     * 1631【并查集】最小体力消耗路径
     * 288 ms
     * 先将每个点与其右边或下边的点连成具有cost的无向边
     * 然后将这些边按cost排序
     * 然后从cost最小的边开始连接两个点到并查集中
     * 当并查集中的起始点与终止点连通起来的时候，最后加入的边的cost即为结果
     */
    public int minimumEffortPath(int[][] heights) {
        if (1 == heights.length && 1 == heights[0].length) {
            return 0;
        }
        //初始化，先将每个点与其右边或下边的点连成具有cost的无向边
        List<Edge1631> edges = new ArrayList<>();
        for (int i=0; i<heights.length; i++) {
            for (int j=0; j<heights[i].length; j++) {
                if (j<heights[i].length-1) {
                    Edge1631 edge = new Edge1631(new Point1631(i, j), new Point1631(i, j + 1), Math.abs(heights[i][j] - heights[i][j + 1]));
                    edges.add(edge);
                }
                if (i<heights.length-1) {
                    Edge1631 edge = new Edge1631(new Point1631(i, j), new Point1631(i+1, j), Math.abs(heights[i][j] - heights[i+1][j]));
                    edges.add(edge);
                }
            }
        }
        //然后将这些边按cost排序
        edges.sort(new Comparator<Edge1631>() {
            @Override
            public int compare(Edge1631 o1, Edge1631 o2) {
                return o1.cost - o2.cost;
            }
        });
        //然后从cost最小的边开始连接两个点到并查集中
        UnionFind1631 uf = new UnionFind1631(new Point1631(0, 0), new Point1631(heights.length-1, heights[heights.length-1].length-1));
        for (Edge1631 edge : edges) {
            if (uf.union(edge.p1, edge.p2)) {
                //当并查集中的起始点与终止点连通起来的时候，最后加入的边的cost即为结果
                return edge.cost;
            }
        }
        return 0;
    }

    static class UnionFind1631 {
        Map<Point1631, Point1631> parent;
        Point1631 start;
        Point1631 end;

        public UnionFind1631(Point1631 start, Point1631 end) {
            parent = new HashMap<>();
            parent.put(start, start);
            parent.put(end, end);
            this.start = start;
            this.end = end;
        }

        public Point1631 find(Point1631 p) {
            if (!parent.containsKey(p)) {
                parent.put(p, p);
                return p;
            }
            Point1631 p_ = parent.get(p);
            if (!p_.equals(parent.get(p_))) {
                parent.put(p, find(p_));
            }
            return parent.get(p);
        }

        public boolean union(Point1631 p1, Point1631 p2) {
            Point1631 root1 = find(p1);
            Point1631 root2 = find(p2);
            if (!root1.equals(root2)) {
                //连接两个点
                parent.put(root1, root2);
            }
            //检查起始点与结束点的连通性
            Point1631 rootStart = find(start);
            Point1631 rootEnd = find(end);
            if (rootStart.equals(rootEnd)) {
                return true;
            }
            return false;
        }
    }

    static class Edge1631 {
        Point1631 p1;
        Point1631 p2;
        int cost;

        public Edge1631(Point1631 p1, Point1631 p2, int cost) {
            this.p1 = p1;
            this.p2 = p2;
            this.cost = cost;
        }
    }

    static class Point1631 {
        int x;
        int y;

        public Point1631(int x, int y) {
            this.x = x;
            this.y = y;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            Point1631 point = (Point1631) o;
            return x == point.x &&
                    y == point.y;
        }

        @Override
        public int hashCode() {
            return Objects.hash(x, y);
        }

        @Override
        public String toString() {
            final StringBuilder sb = new StringBuilder("Point{");
            sb.append("x=").append(x);
            sb.append(", y=").append(y);
            sb.append('}');
            return sb.toString();
        }
    }

    //p1697
    public boolean[] distanceLimitedPathsExist(int n, int[][] edgeList, int[][] queries) {
        boolean[] answer = new boolean[queries.length];
        int[][] queryList = new int[queries.length][4];
        for (int i = 0; i < queries.length; ++i) {
            queryList[i] = new int[]{i, queries[i][0], queries[i][1], queries[i][2]};
        }
        Arrays.sort(queryList, Comparator.comparingInt(q -> q[3]));
        Arrays.sort(edgeList, Comparator.comparingInt(e -> e[2]));
        UnionFind uf = new UnionFind(n);
        int j = 0;
        for (int i = 0; i < queryList.length; ++i) {
            int[] q = queryList[i];
            int limit = q[3];
            while (j < edgeList.length && edgeList[j][2] < limit) {
                uf.union(edgeList[j][0], edgeList[j][1]);
                ++j;
            }
            answer[q[0]] = uf.find(q[1]) == uf.find(q[2]);
        }
        return answer;
    }

    //p1971
    public boolean validPath(int n, int[][] edges, int source, int destination) {
        UnionFind uf = new UnionFind(n);
        for (int[] edge : edges) {
            uf.union(edge[0], edge[1]);
        }
        return uf.find(source) == uf.find(destination);
    }

    public static void main(String[] args) {
        UnionFindSolution solution = new UnionFindSolution();
        Assert.isTrue("[false,true]".equals(JsonUtils.toJson(solution.distanceLimitedPathsExist(3, Parser.parse("[[0,1,2],[1,2,4],[2,0,8],[1,0,16]]"), Parser.parse("[[0,1,2],[0,2,5]]")))));
        Assert.isTrue("[true,false]".equals(JsonUtils.toJson(solution.distanceLimitedPathsExist(5, Parser.parse("[[0,1,10],[1,2,5],[2,3,9],[3,4,13]]"), Parser.parse("[[0,4,14],[1,4,13]]")))));
    }
}
